import numpy as np
import cv2
import os
import argparse
parser = argparse.ArgumentParser(description='Evaluate the saliency map.')
parser.add_argument('--sal_dir', type=str, help='The directory of saliency maps.',default='./results\poolnet_pro-res18\ECSSD_50')
parser.add_argument('--gt_dir', type=str, help='The directory of ground truths.',default='./datasets/sal\ECSSD/test\GT')
THREASHOLDS = 256
EPSILON = 1e-4
BETA = 0.3
def compute_mae(sal, gt):
    mae = 0.0
    height, width = sal.shape[:2]
    mae = np.sum(abs(sal - gt)) / 255.0 / (height * width)
    return mae
def compute_precision_and_recall(sal,gt):
    precision=np.zeros((THREASHOLDS,),dtype=np.float32)
    recall=np.zeros((THREASHOLDS,),dtype=np.float32)
    for th in range(THREASHOLDS):
        sal_th = (sal > th)
        gt_th = (gt > THREASHOLDS / 2)
        tp = np.sum(np.logical_and(sal_th, gt_th))
        a_sum = np.sum(sal_th)
        b_sum = np.sum(gt_th)
        pre = (tp + EPSILON) / (a_sum + EPSILON)
        rec = (tp + EPSILON) / (b_sum + EPSILON)
        precision[th] = pre
        recall[th] = rec
    return precision,recall
def compute_matric(sal_path, gt_path):
    sal = cv2.imread(sal_path, cv2.IMREAD_GRAYSCALE)
    gt = cv2.imread(gt_path, cv2.IMREAD_GRAYSCALE)
    mae = compute_mae(sal, gt)
    precision, recall = compute_precision_and_recall(sal, gt)
    return mae, precision,recall
def do_evaluation(sal_dir, gt_dir):
    img_list = os.listdir(sal_dir)
    img_list = [img.split('.')[0] for img in img_list]
    num_imgs = len(img_list)
    mae = 0.0
    precision = np.zeros((THREASHOLDS,), dtype=np.float32)
    recall = np.zeros((THREASHOLDS,), dtype=np.float32)
    for img_name in img_list:
        sal_path = os.path.join(sal_dir, img_name + '.png')
        gt_path = os.path.join(gt_dir, img_name + '.png')
        _mae, _precision, _recall = compute_matric(sal_path, gt_path)
        mae += _mae
        precision += _precision
        recall += _recall
    precision /= num_imgs
    recall /= num_imgs
    mae /= num_imgs
    fmeasure = (1 + BETA) * precision * recall / (BETA * precision + recall)
    feature_max_index = np.argmax(fmeasure)
    fmeasure_max = fmeasure[feature_max_index]
    print('Max_F-measre:   %.4f' % fmeasure_max)
    print('Mean_F-measre:  %.4f' % np.mean(fmeasure))
    print('Precision:      %.4f' % precision[feature_max_index])
    print('Recall:         %.4f' % recall[feature_max_index])
    print('Mean_Precision: %.4f' % np.mean(precision))
    print('Mean_Recall:    %.4f' % np.mean(recall))
    print('MAE:            %.4f' % mae)
    return fmeasure_max
if __name__ == '__main__':
    args = parser.parse_args()
    do_evaluation(args.sal_dir, args.gt_dir)

    
